#
# CS 143
#
# This file provides the Host class definition.
#
# @author Henry Li (Data Reconstruction)
# @author Alex Fandrianto (Congestion Control)
#

import random
from SimPy.Simulation import *
from Packet import Packet

class Host(Process):
    '''The Host class represents a process which handles packet construction,
    data reconstruction, and congestion control.

    '''
    
    # Metrics metrics
    # What we log to, if it isn't None. Should NEVER BE NONE when running
    metrics = None
    
    __TIMEOUT = 2 # maybe a little small in real life, but...
    
    # These are states for a sending Host.
    __NOT_SENDING = -1
    __SLOW_START = 0
    __CONGESTION_AVOIDANCE = 1
    __FR_FR = 2
    
    # These are the Congestion Control methods we have
    __RENO = 1
    __FAST = 2
    
    __ALPHA = 150 # Try to keep 50 packets in the buffer (when you're the only flow)
    
    controller = None
    
    def __init__(self, nic=None):
        '''Host class constructor.
        
        '''
        # Initialize parent
        Process.__init__(self)
        
        # NIC __nic
        self.__nic = nic

        # dict of pids for outgoing packets with destination ip keys
        # self.__pid_count[destination] = pid
        self.__pid_count = {}

        # dict of sent packets with source ip and pid keys
        # self.__sent_messages[destinaton][pid] = message
        self.__sent_messages = {}

        # dict of received packets with source ip and pid keys
        # self.__data_messages[source][pid] = message
        self.__data_messages = {}
        
        # dict of pids from received packets with source ip and str keys
        # self.__data_indices[source]['start'] = set_of_data_start_pids
        # self.__data_indices[source]['end'] = set_of_data_fin_pids
        # self.__data_indices[source]['received'] = set_of_received_packet_pids
        self.__data_indices = {}

        ######################
        # Congestion Control #
        ######################
        
        # int __state
        # SLOW_START, CONGESTION_AVOIDANCE, or FR_FR
        self.__state = self.__NOT_SENDING
        
        # int __mode
        # FAST or Reno
        # This should only affect Congestion Avoidance's Window Size updates.
        self.__mode = self.__RENO
        
        # float __window_size
        # This must be the case for Congestion_Avoidance.
        self.__window_size = 1.0
        
        # list __send_table of tuples not yet acked
        # [(pid, timeSent, timeoutTime, updateRTTSeed, queueingDelay), ...]
        self.__send_table = []
        
        # int __last_pid_acked
        # Last acked by the receiver. Starts at -1.
        self.__last_pid_acked = -1
        
        # int __num_repeats
        # Starts at 0 for a new ack. Incremented during repeats.
        # It's true that -1 hasn't been repeated at all.
        self.__num_repeats = 0
        
        # bool __resend_FR_FR
        # False until we resend the missing packet for FR FR
        self.__resend_FR_FR = False
        
        # float __eventual_size
        # The window size post-FastRetransmit/FastRecovery.
        # Acts like ssthresh. If -1 OR window_size < this, we're in slow start.
        self.__eventual_size = -1
        
        # float __RTT_min
        # The smallest round trip time ever seen from a newly acked packet.
        # -1 means it hasn't been set yet.
        self.__RTT_min = -1
        
        # float __RTT_last
        # The round trip time of the last new packet acked
        # This estimates the avg RTT. -1 means not set yet.
        self.__RTT_last = -1
        
        # 4-tuple (int, int, int, int) __cur_destination
        # IP address of current destination
        self.__cur_destination = None
        
    def setTimeout(self, timeout):
        self.__TIMEOUT = timeout
    def setALPHA(self, ALPHA):
        self.__ALPHA = ALPHA
    
    
    def go(self):
        ''' Congestion control.
        
        '''
        while True:
            if self.__cur_destination: # I have to send something to someone!
                # first check if anything has timed out.
                if len(self.__send_table) > 0 and self.__send_table[0][2] <= now(): # Dangit, a timeout!!!
                    self.__eventual_size = self.__window_size / 2 # The new ssthresh
                    self.__window_size = 1.0
                    self.__send_table= []
                    pid_to_send = self.__last_pid_acked + 1
                    
                    print "We got a timeout. Reset stuff. T_T. Resending %d at %f" % (pid_to_send, now())
                    
                    packet = self.__sent_messages[self.__cur_destination][pid_to_send]
                    self.__send_table.append((packet.pid, now(), now() + self.__TIMEOUT, packet.updateRTTSeed, self.__nic.getQueueingDelay()))
                    
                    # Please remind me to check for timeout at this time
                    self.controller.remindMe(self, now() + self.__TIMEOUT)
                    self.__nic.send(packet)
            
                pid_to_send = self.__last_pid_acked + 1
                    
                # Handle the special Fast Retransmit case
                if self.__state == self.__FR_FR and self.__resend_FR_FR == False:
                    packet = self.__sent_messages[self.__cur_destination][pid_to_send]
                    # Note that packet.pid = pid_to_send
                    self.__send_table.append((packet.pid, now(), now() + self.__TIMEOUT, packet.updateRTTSeed, self.__nic.getQueueingDelay()))
                    # It's already in the table.
                    
                    # Please remind me to check for timeout at this time
                    self.controller.remindMe(self, now() + self.__TIMEOUT)
                    print "Resending %d with window %f and state 2 FR_FR at %f" % (pid_to_send, self.__window_size, now())
                    self.__nic.send(packet)
                    
                    self.__resend_FR_FR = True
                    
                # look through the send table and up the newest packet to send
                for i in range(len(self.__send_table)):
                    if self.__send_table[i][0] == pid_to_send: # packet already here
                        pid_to_send += 1
                        
                while len(self.__send_table) < int(self.__window_size) and pid_to_send <= max(self.__sent_messages[self.__cur_destination]):
                    packet = self.__sent_messages[self.__cur_destination][pid_to_send]
                    # Note that packet.pid = pid_to_send
                    self.__send_table.append((packet.pid, now(), now() + self.__TIMEOUT, packet.updateRTTSeed, self.__nic.getQueueingDelay()))
                    
                    # Please remind me to check for timeout at this time
                    self.controller.remindMe(self, now() + self.__TIMEOUT)
                    print "Sending %d with window %f and state %d at %f" % (pid_to_send, self.__window_size, self.__state, now())
                    self.__nic.send(packet)
                    pid_to_send += 1
            yield passivate, self
            
    def addNIC(self, nic):
        '''Set the NIC object.
        
        Input:
            nic: NIC
        '''
        self.__nic = nic
        
    def send(self, data, packet_size, destination, ccType):
        '''Split data into packets and give to congestion control.
        
        Input:
            data: str
            packet_size: int
            destination: 4 tuple representing an IP address
            ccType: str (either FAST or Reno)

        '''
        # Prevent misuse when host's nic has not been set.
        if not self.__nic:
            print 'Error: Host has no nic'
            return None
    
        # Initialize `__pid_count` for specific `destinaton`
        if destination not in self.__pid_count:
            self.__pid_count[destination] = 0
        
        # Initialize `__sent_messages` for specific `destinaton`
        if destination not in self.__sent_messages:
            self.__sent_messages[destination] = {}
        
        # Split data into list of packets
        messages = self.__split_data(data, packet_size)
        
        # Index of last packet
        last_message = len(messages) - 1
        
        # For congestion control...
        self.__state = self.__SLOW_START
        if ccType == 'FAST':
            self.__mode = self.__FAST
        else:
            self.__mode = self.__RENO
        self.__window_size = 1.0
        self.__send_table = []
        self.__last_pid_acked = self.__pid_count[destination] - 1
        self.__num_repeats = 0
        self.__resend_FR_FR = False
        self.__eventual_size = -1
        self.__RTT_min = -1
        self.__RTT_last = -1
        self.__cur_destination = destination
        
        for index, message in enumerate(messages):
            # Create packets where last packet has `fin` marker
            fin = True if index == last_message else False
            packet = self.__create_packet(destination, packet_size, 
                                          message, fin)
            
            # Store packet with destinaton and pid                          
            self.__sent_messages[destination][packet.pid] = packet
            
            # Update `__pid_count` for specific `destination`
            self.__pid_count[destination] += 1
            

        # Reactivate in case we passivated
        if self.passive():
            reactivate(self)
            
    
    def receive(self, nic, packet):
        '''
        nic: NIC object
        packet: Packet object
        
        Send acknowledgment to packet's source.
        
        Return a data string if all packets in the data sequence have been
        received, else store packets for future reconstruction.
        '''
        # Prevent misuse when host's nic has not been set.
        if not self.__nic:
            print "Error: Host has no nic"
            return None
            
        # Reactivate in case we passivated
        if self.passive():
            reactivate(self)
            
        # Used to log the old window size
        windowTemp = self.__window_size
            
        # Check if packet is intended for this host
        if packet.destination == self.__nic.getIP():

            if packet.ack:
                print "Got ack for %d with last best %d" % (packet.pid, self.__last_pid_acked)
                if self.__last_pid_acked == packet.pid: # repeat ack!
                    self.__num_repeats += 1
                    
                    # Switch states if you see too many repeats
                    if self.__state == self.__CONGESTION_AVOIDANCE and self.__num_repeats == 3:
                        print "I entered FR_FR"
                        self.__state = self.__FR_FR
                        self.__eventual_size = max(self.__window_size / 2, 1.0)
                        self.__window_size = self.__eventual_size + self.__num_repeats
                    elif self.__state == self.__FR_FR: # handle window size increase here; Fast Recovery
                        self.__window_size = self.__eventual_size + self.__num_repeats
                        
                elif packet.pid < self.__last_pid_acked: # We got this already, pass
                    pass
                elif packet.pid > self.__last_pid_acked: # It's a different ack
                    self.__last_pid_acked = packet.pid
                    self.__num_repeats = 0
                    self.__resend_FR_FR = False
                    
                    if self.__state == self.__FR_FR:
                        # time to leave __FR_FR
                        print "I left FR_FR"
                        self.__state = self.__CONGESTION_AVOIDANCE
                        self.__window_size = self.__eventual_size
                    
                    # I need to find the RTT. This only matters for FAST
                    for x in self.__send_table:
                        if x[0] == self.__last_pid_acked:
                            if packet.updateRTTSeed == x[3]:
                                self.__RTT_last = now() - x[1]
                                self.metrics.log(self, now(), packet, 'RTT', [self.__RTT_last, windowTemp, self.__last_pid_acked]) 
                            
                                # Adjust by current queueing delay and old queueing delay since this is the Estimated RTT
                                self.__RTT_last += self.__nic.getQueueingDelay() - x[4]
                            
                                # Update the RTT min, if necessary
                                if self.__RTT_min == -1 or self.__RTT_min >= self.__RTT_last:
                                    self.__RTT_min = self.__RTT_last
                                
                                print 'Ack %d at %f. Cur RTT %f and Min RTT %f' % (self.__last_pid_acked, now(), self.__RTT_last, self.__RTT_min)
                            else:
                                print 'Ack %d at %f. but is not for this entry; this is okay' % (self.__last_pid_acked, now())
                    
                    # Now remove from the table any entries whose pid <= this. (x[0] is a pid)
                    self.__send_table = [x for x in self.__send_table if not x[0] <= self.__last_pid_acked]    

                    # Adjust window size based on state and CC mode
                    if self.__state == self.__SLOW_START:
                        self.__window_size += 1.0
                        
                        # Recall that after the first timeout, regrowth to ssthresh puts you at __CONGESTION_AVOIDANCE
                        if self.__eventual_size != -1 and self.__window_size > self.__eventual_size:
                            self.__state = self.__CONGESTION_AVOIDANCE
                    elif self.__state == self.__CONGESTION_AVOIDANCE and self.__mode == self.__RENO:
                        self.__window_size += 1./self.__window_size
                    elif self.__state == self.__CONGESTION_AVOIDANCE and self.__mode == self.__FAST:
                        # Instead of updating the window all at once, we can gradually change the window size.
                        # Growth will be exponential this way.
                        # self.__window_size = self.__window_size * self.__RTT_min / self.__RTT_last + self.__ALPHA
                        if self.__window_size > self.__window_size * self.__RTT_min / self.__RTT_last + self.__ALPHA:
                            self.__window_size -= 1
                        elif self.__window_size < self.__window_size * self.__RTT_min / self.__RTT_last + self.__ALPHA:
                            self.__window_size += 1
                        
                
                # But wait! Are we done? If so, reset everything!
                if self.__cur_destination != None and self.__last_pid_acked == max(self.__sent_messages[self.__cur_destination]):
                    self.__state = self.__NOT_SENDING
                    self.__mode = self.__RENO
                    self.__window_size = 1.0
                    self.__send_table = []
                    self.__last_pid_acked = -1
                    self.__num_repeats = 0 # It's true that -1 hasn't been repeated at all.
                    self.__resend_FR_FR = False # True after the missing packet was resent for FR_FR
                    self.__eventual_size = -1 # This acts like ssthresh
                    self.__RTT_min = -1 # Not set yet
                    self.__RTT_last = -1 # Not set yet
                    self.__cur_destination = None
                    print "We're done sending this flow!"
                
                # Reactivate in case we passivated
                if self.passive():
                    reactivate(self)
                
                return None # No more to be done since this was an ack

            # Initialize `__data_messages` for specific source
            if packet.source not in self.__data_messages:
                self.__data_messages[packet.source] = {}
            
            # Initialize `__data_indices` for specific source
            if packet.source not in self.__data_indices:
                self.__data_indices[packet.source] = {'start': set([0]),
                                                      'end': set(),
                                                      'received': set()}
            
            # Access the index sets in `__data_indices` for specific source
            start_indices = self.__data_indices[packet.source]['start']
            end_indices = self.__data_indices[packet.source]['end']
            received_packets = self.__data_indices[packet.source]['received']
            
            # Start index of earliest data sequence
            min_start = min(start_indices)
            
            # Received this packet before so do nothing
            if (packet.pid < min_start) or (packet.pid in received_packets):
                return None

            # Add current packet to set of received packets
            received_packets.add(packet.pid)

            # Log that new packet has been received
            self.metrics.log(self, now(), packet, 'NewPacketReceived',
                             len(received_packets))
            
            # Store message with source and pid
            self.__data_messages[packet.source][packet.pid] = packet.message

            # Send ack back to received packet's source
            acked = self.__send_ack(min_start, received_packets, packet.source,
                                    packet.updateRTTSeed)
            print "Got packet %d and acking for %d at %f" % (packet.pid, acked, 
                                                             now())

            # Update start and end index if receiving last packet in data block
            if packet.fin:
                print "Yay, we got fin!", packet.fin
                end_indices.add(packet.pid)
                start_indices.add(packet.pid + 1)
                self.metrics.log(self, now(), None, 'FlowEnd')
            
            # Check for complete data only if a packet with `fin` marker exists
            if end_indices != set():
                min_end = min(end_indices)
                
                # Range of pids for the earliest block of data
                data_range = range(min_start, min_end + 1)
                data_chunk = set(data_range)
                
                # Data is complete only if the data chunk has been received
                if data_chunk - received_packets == set():
                    
                    # Delete all indices related to data chunk because we got it
                    received_packets -= data_chunk
                    start_indices.remove(min_start)
                    end_indices.remove(min_end)
                    
                    # Reconstruct the data and return
                    complete_data = ''
                    for pid in data_range:
                        complete_data += \
                                        self.__data_messages[packet.source][pid]
                        del self.__data_messages[packet.source][pid]
                    print "We're done receiving! I got dis many %d" \
                          % len(complete_data)
                    return complete_data
            
    def disconnect(self, nic):
        '''Host is at one end of the network, so disconnecting does nothing.
        This function definition is to allow NIC to call disconnect on its
        parent device.
        
        Input:
            nic: NIC

        '''
        pass

    def __split_data(self, seq, length):
        '''Divide the sequence into list of subsequences of the given length and
        return this list.
        
        Input:
            seq: str or any sequence
            length: int

        Output:
            subsequences: list

        '''
        return [seq[i:i+length] for i in range(0, len(seq), length)]
        
    def __create_packet(self, destination, packet_size, message, fin=False):
        '''Create a packet to send.
        
        Input:
            destination: 4 tuple representing an IP address
            packet_size: int
            message: str
            fin: bool
        
        Output:
            packet: Packet

        '''
        packet = Packet()
        packet.pid = self.__pid_count[destination]
        packet.source = self.__nic.getIP()
        packet.destination = destination
        packet.size = packet_size
        packet.message = message
        packet.fin = fin
        packet.updateRTTSeed = random.random()
        return packet
    
    def __create_ack(self, pid, destination, updateRTTSeed):
        '''Create an acknowledgment packet.
        
        Input:
            pid: int
            destination: 4 tuple representing an IP address
            updateRTTSeed: int
        
        Output:
            packet: Packet

        '''
        ack_packet = Packet()
        ack_packet.ack = True
        ack_packet.size = 1000
        ack_packet.source = self.__nic.getIP()
        ack_packet.destination = destination
        ack_packet.pid = pid
        ack_packet.updateRTTSeed = updateRTTSeed
        return ack_packet
    
    def __send_ack(self, min_start, received_packets, source, updateRTTSeed):
        '''Send acknowledgment packet for largest pid received or for 
        last packet received before a missing packet.
        
        Input:
            min_start: int
            received_packets: set
            source: 4 tuple representing an IP address
            updateRTTSeed: int
        
        Output:
            pid: int

        '''
        # Determine the contiguous block of packets that we should have
        # and get set of missing packets if any
        max_received = max(received_packets)
        expected_packets = set(range(min_start, max_received + 1))
        missing_packets = expected_packets - received_packets

        if missing_packets:
            # Create acknowledgment for the last received packet before 
            # the first missing packet, or pid = -1 if first packet lost
            ack_packet = self.__create_ack(min(missing_packets) - 1, source, 0)
        else:
            # Create acknowledgment for packet with highest received pid
            # because no packets are missing
            ack_packet = self.__create_ack(max_received, source, updateRTTSeed)
        
        # Send acks through nic directly
        self.__nic.send(ack_packet)
        
        return ack_packet.pid
